# coding: utf-8
from scipy.stats import rankdata
import copy
from ..op_basic import as_op
import pandas as pd


def __cal_group_rank(values, rank_method):
    '''
    内部函数
    '''
    return rankdata(values.values, method=rank_method)


def __cal_group_rank_reverse(values, rank_method):
    '''
    内部函数
    '''
    return rankdata(-values.values, method=rank_method)


def __cal_group_rank_norm(values, rank_method):
    '''
    内部函数
    '''
    if len(values.values) == 1:
        return 0
    else:
        return (rankdata(values.values, method=rank_method) - 1) / (len(values.values) - 1)  # 放缩到0~1


def __cal_group_rank_reverse_norm(values, rank_method):
    '''
    内部函数
    '''
    if len(values.values) == 1:
        return 0
    else:
        return (rankdata(-values.values, method=rank_method) - 1) / (len(values.values) - 1)  # 放缩到0~1


@as_op('group_rank')
def group_rank(df, group_columns, rank_columns, prefix='rank', return_way='new', rank_method='min',
               gen_origin=True, gen_reverse=False, gen_norm=False, gen_reverse_norm=False):
    '''
    对df计算组内rank。

    要考虑两个方面：
        对于空值(尝试哪种方法更好)：1、先fill，再处理 2、按照这里的方法组类fill 
        对于相同值：参考不同的method
        生成怎样的排名：是从大到小，还是从小到大，要不要组内缩放到0~1

    输入：
        df：要求的df。
        group_columns：按这些列分组
        rank_columns：计算这些列的rank，一般都是些连续变量或组合后的连续变量
        prefix：计算后列的prefix
        return_way：new表示返回新的同index的df，add表示直接concat。
        rank_method：对相同的value的rank如何处理。参考https://docs.scipy.org/doc/scipy-0.16.0/reference/generated/scipy.stats.rankdata.html

        gen_origin: 是否生成从小到大的排名，这是默认的。默认为True，可以置为False不生成。
        gen_reverse：是否生成从大到小的排名。排名的列名为prefix +'_' + 原列名 + '_reverse'
        gen_norm: 是否生成从小到大标准化（同一组映射到0~1）的排名。排名的列名为prefix +'_' + 原列名 + '_norm'
        gen_reverse_norm：是否生成从大到小标准化的排名。排名的列名为prefix +'_' + 原列名 + '_reverse_norm'

    输出：
        新的df或concat的df。
    '''
    assert (df.columns.duplicated() == False).all()
    for col in group_columns:
        assert col in df.columns, col
        assert col not in rank_columns, col
    for col in rank_columns:
        assert col in df.columns, col
        assert col not in group_columns, col
    assert gen_origin or gen_reverse or gen_norm or gen_reverse_norm

    consider_col = copy.copy(group_columns)
    consider_col.extend(rank_columns)
    consider_df = df.filter(consider_col, axis=1)  # 这里创建一个新的df

    # 处理空值。分组按排名填充
    for col in rank_columns:
        if consider_df[col].notnull().all() == False:
            t_col = copy.copy(group_columns)
            t_col.append(col)
            consider_df[col] = consider_df[t_col].groupby(group_columns).transform(lambda x: x.fillna(x.median()))  # 感觉用中位数填充更加科学
            consider_df[col].fillna(consider_df[col].median(), inplace=True)   # 分组后全为nan的现在还是nan，所以要fill。实际上这里填充什么都无所谓，因为后面在组内都是一样的排名

    result = []
    grouped = consider_df.groupby(group_columns)

    if gen_origin:
        __result = grouped.transform(__cal_group_rank, rank_method=rank_method)
        __origin_result = __result
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col
        result.append(__result.rename(columns=__rename_dict))

    if gen_reverse:  # 这里貌似不可以通过先求组长度再减去原来排名的方式来缩短时间。如min就可以举出反例。
        __result = grouped.transform(__cal_group_rank_reverse, rank_method=rank_method)
        __reverse_result = __result
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_reverse'
        result.append(__result.rename(columns=__rename_dict))

    if gen_norm:
        if gen_origin:  # 这样不重复计算
            grouplen = grouped.transform(lambda x: len(x))
            grouplen = (grouplen - 1).replace(0, 1)  # 防止除以0
            __result = (__origin_result - 1) / grouplen
        else:
            __result = grouped.transform(__cal_group_rank_norm, rank_method=rank_method)
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_norm'
        __result.rename(columns=__rename_dict, inplace=True)
        result.append(__result)

    if gen_reverse_norm:
        if gen_reverse:  # 这样不重复计算
            grouplen = grouped.transform(lambda x: len(x))
            grouplen = (grouplen - 1).replace(0, 1)  # 防止除以0
            __result = (__reverse_result - 1) / grouplen
        else:
            __result = grouped.transform(__cal_group_rank_reverse_norm, rank_method=rank_method)
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_reverse_norm'
        __result.rename(columns=__rename_dict, inplace=True)
        result.append(__result)

    result = pd.concat(result, axis=1, copy=False)

    if return_way == 'new':
        return result
    elif return_way == 'add':
        return pd.concat([df, result], axis=1, copy=False)
    else:
        raise Exception("unknown return_way")

@as_op('group_rank_by_other')
def group_rank_by_other(df, other_list, rank_columns, prefix='rank', return_way='new', rank_method='min',
               gen_origin=True, gen_reverse=False, gen_norm=False, gen_reverse_norm=False):
    '''
    对df计算组内rank。此时的group标准是其他df，如[ids.a, ids.b]

    要考虑两个方面：
        对于空值(尝试哪种方法更好)：1、先fill，再处理 2、按照这里的方法组类fill 
        对于相同值：参考不同的method
        生成怎样的排名：是从大到小，还是从小到大，要不要组内缩放到0~1

    输入：k
        df：要求的df。
        other_list：按这些列分组
        rank_columns：计算这些列的rank，一般都是些连续变量或组合后的连续变量
        prefix：计算后列的prefix
        return_way：new表示返回新的同index的df，add表示直接concat。
        rank_method：对相同的value的rank如何处理。参考https://docs.scipy.org/doc/scipy-0.16.0/reference/generated/scipy.stats.rankdata.html

        gen_origin: 是否生成从小到大的排名，这是默认的。默认为True，可以置为False不生成。
        gen_reverse：是否生成从大到小的排名。排名的列名为prefix +'_' + 原列名 + '_reverse'
        gen_norm: 是否生成从小到大标准化（同一组映射到0~1）的排名。排名的列名为prefix +'_' + 原列名 + '_norm'
        gen_reverse_norm：是否生成从大到小标准化的排名。排名的列名为prefix +'_' + 原列名 + '_reverse_norm'

    输出：
        新的df或concat的df。
    '''
    assert (df.columns.duplicated() == False).all()
    for col in rank_columns:
        assert col in df.columns, col
    assert gen_origin or gen_reverse or gen_norm or gen_reverse_norm

    consider_col = rank_columns
    consider_df = df.filter(consider_col, axis=1)  # 这里创建一个新的df

    # 处理空值。分组按排名填充
    for col in rank_columns:
        if consider_df[col].notnull().all() == False:
            consider_df[col] = consider_df[col].groupby(other_list).transform(lambda x: x.fillna(x.median()))  # 感觉用中位数填充更加科学
            consider_df[col].fillna(consider_df[col].median(), inplace=True)   # 分组后全为nan的现在还是nan，所以要fill。实际上这里填充什么都无所谓，因为后面在组内都是一样的排名

    result = []
    grouped = consider_df.groupby(other_list)

    if gen_origin:
        __result = grouped.transform(__cal_group_rank, rank_method=rank_method)
        __origin_result = __result
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col
        result.append(__result.rename(columns=__rename_dict))

    if gen_reverse:  # 这里貌似不可以通过先求组长度再减去原来排名的方式来缩短时间。如min就可以举出反例。
        __result = grouped.transform(__cal_group_rank_reverse, rank_method=rank_method)
        __reverse_result = __result
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_reverse'
        result.append(__result.rename(columns=__rename_dict))

    if gen_norm:
        if gen_origin:  # 这样不重复计算
            grouplen = grouped.transform(lambda x: len(x))
            grouplen = (grouplen - 1).replace(0, 1)  # 防止除以0
            __result = (__origin_result - 1) / grouplen
        else:
            __result = grouped.transform(__cal_group_rank_norm, rank_method=rank_method)
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_norm'
        __result.rename(columns=__rename_dict, inplace=True)
        result.append(__result)

    if gen_reverse_norm:
        if gen_reverse:  # 这样不重复计算
            grouplen = grouped.transform(lambda x: len(x))
            grouplen = (grouplen - 1).replace(0, 1)  # 防止除以0
            __result = (__reverse_result - 1) / grouplen
        else:
            __result = grouped.transform(__cal_group_rank_reverse_norm, rank_method=rank_method)
        __rename_dict = {}
        for col in rank_columns:
            __rename_dict[col] = prefix + '_' + col + '_reverse_norm'
        __result.rename(columns=__rename_dict, inplace=True)
        result.append(__result)

    result = pd.concat(result, axis=1, copy=False)

    if return_way == 'new':
        return result
    elif return_way == 'add':
        return pd.concat([df, result], axis=1, copy=False)
    else:
        raise Exception("unknown return_way")
if __name__ == '__main__':
    df = pd.DataFrame({'a': [1, 1, 1, 2, 2, 3, 3, 3, 3], 'b': [0.1, 0.8, 0.4, 0.5, 0.4, 0.1, 0.2, 0.2, 0.1]})
    group_rank(df, ['a'], ['b'], return_way='add', gen_reverse=True, rank_method='average', gen_norm=True, gen_reverse_norm=True)
